package com.example.langchaindemo.base;

import com.example.langchaindemo.config.Langchain4jProperties;
import com.example.langchaindemo.service.AiAssistant;
import com.example.langchaindemo.utils4llm.Calculator;
import com.google.gson.JsonElement;
import com.google.gson.JsonObject;
import com.google.gson.JsonParser;
import dev.ai4j.openai4j.chat.ToolMessage;
import dev.langchain4j.agent.tool.ToolExecutionRequest;
import dev.langchain4j.agent.tool.ToolSpecification;
import dev.langchain4j.agent.tool.ToolSpecifications;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ToolExecutionResultMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.request.ChatRequest;
import dev.langchain4j.model.chat.request.DefaultChatRequestParameters;
import dev.langchain4j.model.chat.response.ChatResponse;
import lombok.RequiredArgsConstructor;
import org.springframework.util.CollectionUtils;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.util.List;

/**
 * funcation call调用
 */
@RequestMapping("/api")
@RestController
@RequiredArgsConstructor
public class ChatWithTools {

    final ChatLanguageModel model;
    final AiAssistant assistant;

    final Langchain4jProperties properties;

    @GetMapping("/lowlevel/tool")
    public String lowChat(@RequestParam("message") String message) {
        List<ToolSpecification> toolSpecifications = ToolSpecifications.toolSpecificationsFrom(Calculator.class);
        UserMessage userMessage = UserMessage.from(message);
        ChatResponse chat = model.chat(ChatRequest.builder().messages(userMessage).parameters(DefaultChatRequestParameters.builder().toolSpecifications(toolSpecifications).build()).build());
        MessageWindowChatMemory memory = new MessageWindowChatMemory.Builder().maxMessages(20).build();
        memory.add(userMessage);
        memory.add(chat.aiMessage());

        List<ToolExecutionRequest> toolExecutionRequests = chat.aiMessage().toolExecutionRequests();
        if (!CollectionUtils.isEmpty(toolExecutionRequests)) {
            for (ToolExecutionRequest tool : toolExecutionRequests) {
                String arguments = tool.arguments();
                String name = tool.name();
                System.out.println(arguments + "\n" + name);
                String response = switch (name) {
                    case "stringLength" -> new Calculator().stringLength(arguments) + "";
                    case "sqrt" -> new Calculator().sqrt(Integer.valueOf(arguments)) + "";
                    default -> "unknown tool";
                };
                if (name.equalsIgnoreCase("add")) {
                    JsonObject jobj = JsonParser.parseString(arguments).getAsJsonObject();
                    response = new Calculator()
                            .add(Integer.valueOf(jobj.get("a").toString()), Integer.valueOf(jobj.get("b").toString())) + "";
                }

                ToolExecutionResultMessage toolMessage = ToolExecutionResultMessage.from(tool, response);
                memory.add(toolMessage);
            }
        }
        ChatResponse finalResult = model.chat(memory.messages());
        memory.add(finalResult.aiMessage());
        return finalResult.aiMessage().text();
    }

    @GetMapping("/highlevel/tool")
    public String highChat(@RequestParam("message") String message) {
        String answer = assistant.chat(message);
        return answer;
    }
}
